(* "Auxiliary" functions. Used by the main checker functions. *)

open Spitree

module SSet = Set.Make(String);;

(* Name substitution for dependent pair-types
 * We assume types from an alpha-converted process.
 * Input:
 * typ typ - Type to perform substitution in
 * name current - Name to substitute
 * msg newmsg - Message to substitute name with
 * Output: typ - Substituted type
 *)
let rec subst_name typ current newmsg =
  let rec in_msg msg =
    match msg with
      | Name(x) ->
	  if x = current then
	    newmsg
	  else
	    msg
      | Encr(m, n) ->
	  Encr(in_msg m, in_msg n)
      | MPair(m, n) ->
	  MPair(in_msg m, in_msg n)
      | Ok -> Ok
  in
  let in_literal (pred, termlist) =
    let subst_termlist =
      List.map (fun term ->
		  match term with
		    | Var(x) -> term
		    | SpiMsg(msg) -> SpiMsg(in_msg msg)
	       ) termlist
    in
      (pred, subst_termlist)
  in
  let in_hclause clause =
    List.map (fun literal ->
		in_literal literal
	     ) clause
  in
  let in_data data =
    List.map (fun hclause ->
		in_hclause hclause
	     ) data
(*  let rec in_reso reso =
    match reso with
      | Reso(l, m) -> Reso(l, in_msg m)
      | RPair(r, s) -> RPair(in_reso r, in_reso r)
      | Empty -> Empty*)
  in
  match typ with
    | Un -> Un
    | Key(t) -> Key(subst_name t current newmsg)
    | Ch(t) -> Ch(subst_name t current newmsg)
    | TPair(x, t, u) -> TPair(x, subst_name t current newmsg,
			      subst_name u current newmsg)
    | TOk(d) -> TOk(in_data d)
    | Unknown -> Unknown;;

(* Decide type equality. Dependent pairs are taken into account using subst_name
 * Input:
 * int acc - Accumulator for name substitution in dependent pair types.
 * typ t1, t2 - Types for comparison.
 * Output: bool, indicating equality between t1 and t2 *)
(*exception TypesNotEqual of typ*typ;;*)
let rec type_equal ?(acc=0) t1 t2 =
  match (t1, t2) with
    | (Un, Un) -> true
    | (Key(t), Key(u)) -> (type_equal ~acc:acc t u)
    | (Ch(t), Ch(u)) -> (type_equal ~acc:acc t u)
    | (TPair(x, left1, right1), TPair(y, left2, right2)) ->
	if (type_equal ~acc:acc left1 left2) then
	  let r1_subst = subst_name right1 x (Name(string_of_int acc)) in
	  let r2_subst = subst_name right2 y (Name(string_of_int acc)) in
	    type_equal ~acc:(acc + 1) r1_subst r2_subst
	else false (*raise (TypesNotEqual (t1, t2))*)
    | (TOk(d1), TOk(d2)) ->
	(List.for_all (fun d1elt ->
			 List.mem d1elt d2
		      ) d1) &&
	  (List.for_all (fun d2elt ->
			   List.mem d2elt d1
			) d2)
    | _ -> false;;

(* Find the domain of an environment
 * Input: list env_element env - Environment
 * Output: String Set - A set of strings representing the domain of env
 * Note on list-folding:
 * The reason we use fold_left here, is
 * because fold_left is tail-recursive, and because we are
 * folding into a set, order becomes irrelevant.
 * fold_right is used in the effects function to preserve the
 * order when folding into a new list.
 *)
let dom env =
  List.fold_left (fun sset typ ->
		    match typ with
		      | Type(x, _) -> SSet.add x sset
		      | _ -> sset) SSet.empty env;;

(* Flatten effect to a list. Makes it easier to test for
 * effect inclusion and extract effects.
 * Input: reso s - Effect
 * Output: list reso - Flattened list of effects (type Reso)
 *)
(*let rec effect_to_list s =
  match s with
    | Reso(_, _) -> [s]
    | RPair(r, t) ->
	List.append (effect_to_list r) (effect_to_list t)
    | Empty -> [];;*)

(* Extract effects from environment
 * Input: list env_element env - Environment
 * Output: list reso - List of effects from environment
 *)
(*let effects env =
  List.fold_right (fun eff list ->
		     match eff with
		       | Effect(r) -> r::list
		       | Type(_, TOk(s)) -> List.append (effect_to_list s) list
		       | _ -> list
		  ) env [];;*)

let clauses env =
  List.fold_right (fun eff list ->
		     match eff with
		       | Effect(d) -> List.append d list
		       | Type(_, TOk(d)) -> List.append d list
		       | _ -> list
		  ) env [];;

(* Generative type?
 * Input: typ typ - Type
 * Output: Boolean indicating if type typ is generative *)
exception NotGenerative of typ;;
let generative typ =
  match typ with
    | Un -> true
    | Key(_) -> true
    | Ch(_) -> true
    | _ -> false;;

(* Extract environment from a process. Corresponds to definition of env function
 * Input: proc proc - Process
 * Output: Environment of process proc *)
let rec env_of proc =
  match proc with
    | PPara(p, q) -> List.append (env_of q) (env_of p)
    | Nu(x, t, p) -> List.append (env_of p) [Type(x, t)]
    | Datprog(d) -> [Effect(d)]
(*    | Begin(l, m) -> [Effect(Reso(l, m))]*)
    | _ -> [];;

(* Free name functions
 * fn_msg, fn_res, fn_typ:
 * Input: msg msg, reso res, typ typ, respectively.
 * Output: String Set - Set of strings indicating the free names of the input. *)
let rec fn_msg msg =
  match msg with
    | Name(x) ->
	if x = "*" then
	  SSet.empty
	else
	  SSet.add x SSet.empty
    | Encr(m, n) -> SSet.union (fn_msg m) (fn_msg n)
    | MPair(m, n) -> SSet.union (fn_msg m) (fn_msg n)
    | _ -> SSet.empty;;
(*let rec fn_res res =
  match res with
    | Reso(_, m) -> fn_msg m
    | RPair(r, s) -> SSet.union (fn_res r) (fn_res s)
    | _ -> SSet.empty;;*)
let fn_literal (_, termlist) =
  List.fold_left (fun accum_set term ->
		    match term with
		      | SpiMsg(msg) ->
			  SSet.union accum_set (fn_msg msg)
		      | _ -> accum_set
		 ) SSet.empty termlist;;
let fn_hclause clause =
  List.fold_left (fun accum_set literal ->
		    SSet.union accum_set (fn_literal literal)
		 ) SSet.empty clause;;
let fn_data data =
  List.fold_left (fun accum_set clause ->
		    SSet.union accum_set (fn_hclause clause)
		 ) SSet.empty data;;
let rec fn_typ typ =
  match typ with
    | Key(t) -> fn_typ t
    | Ch(t) -> fn_typ t
    | TPair(x, t, u) ->
	SSet.union (fn_typ t) (fn_typ (subst_name u x (Name("*"))))
    | TOk(d) -> fn_data d
    | _ -> SSet.empty;;

(* Effect Inclusion
 * Input: list Reso - List of effects. Important: Flatten before passing.
 * Output: bool, indicating if all effects in s1 is included in s2.
 *)
exception EffectInclusion of string;;
let rec effectinclusion s1 s2 =
  match s1 with
    | [] -> true
    | (x::effectrest) ->
	if (List.mem x s2) then
	  effectinclusion effectrest s2
	else
	  raise (EffectInclusion ("Effect inclusion failed: " ^ (string_of_reso x) ^ " not permitted"));;

(* Name extraction
 * Input:
 * name x - Name
 * list env_element env - Environment
 * Output: typ - Indicating the type of x in environment env
 *)
exception NoName of string;;
let rec nametype x env =
  match env with
    | [] -> raise (NoName ("Name " ^ x ^ " not found in environment"))
    | Type(name, typ)::restenv ->
	if x = name then
	  typ
	else
	  nametype x restenv
    | Effect(_)::restenv ->
	nametype x restenv;;
